#!/usr/bin/env python3
"""
数据预处理模块
用于加载和处理LeapGestRecog手势数据集
"""

import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
from tqdm import tqdm

def create_dataframe_from_directory(data_dir="data/leapGestRecog"):
    """
    扫描原始数据目录，创建 DataFrame（不移动文件）
    返回包含文件路径和标签的 DataFrame
    """
    import pandas as pd
    
    print("正在扫描数据目录...")
    
    if not os.path.exists(data_dir):
        print(f"错误：数据目录 '{data_dir}' 不存在")
        return None, None
    
    file_paths = []
    labels = []
    class_names = []
    
    # 遍历用户目录
    for user_dir in os.listdir(data_dir):
        user_path = os.path.join(data_dir, user_dir)
        if not os.path.isdir(user_path) or not user_dir.isdigit():
            continue
        
        # 遍历手势目录
        for gesture_dir in os.listdir(user_path):
            gesture_path = os.path.join(user_path, gesture_dir)
            if not os.path.isdir(gesture_path):
                continue
            
            # 提取手势名称
            gesture_name = gesture_dir.split('_', 1)[1] if '_' in gesture_dir else gesture_dir
            
            if gesture_name not in class_names:
                class_names.append(gesture_name)
            
            # 收集图像文件路径
            for img_file in os.listdir(gesture_path):
                if img_file.lower().endswith('.png'):
                    img_path = os.path.join(gesture_path, img_file)
                    file_paths.append(img_path)
                    labels.append(gesture_name)
    
    # 创建 DataFrame
    df = pd.DataFrame({
        'filename': file_paths,
        'class': labels
    })
    
    print(f"✓ 扫描完成:")
    print(f"  - 总图像数: {len(df)}")
    print(f"  - 手势类别: {sorted(class_names)}")
    
    # 统计每个类别的数量
    for class_name in sorted(class_names):
        count = len(df[df['class'] == class_name])
        print(f"  - {class_name}: {count} 张")
    
    return df, sorted(class_names)

def create_data_generators_from_dataframe(df, class_names, batch_size=32, validation_split=0.2):
    """
    从 DataFrame 创建数据生成器（直接从原始位置读取，不移动文件）
    """
    from sklearn.model_selection import train_test_split
    
    print("正在创建数据生成器（从原始位置读取）...")
    
    # 分割训练集和验证集
    train_df, val_df = train_test_split(
        df, 
        test_size=validation_split, 
        stratify=df['class'],
        random_state=42
    )
    
    print(f"数据分割完成：")
    print(f"- 训练集: {len(train_df)} 样本")
    print(f"- 验证集: {len(val_df)} 样本")
    print(f"- 批次大小: {batch_size}")
    
    # 训练数据生成器（带数据增强）
    train_datagen = ImageDataGenerator(
        rescale=1./255,
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )
    
    # 验证数据生成器（不增强）
    val_datagen = ImageDataGenerator(rescale=1./255)
    
    # 训练集生成器
    train_generator = train_datagen.flow_from_dataframe(
        train_df,
        x_col='filename',
        y_col='class',
        target_size=(224, 224),
        batch_size=batch_size,
        class_mode='categorical',
        shuffle=True
    )
    
    # 验证集生成器
    val_generator = val_datagen.flow_from_dataframe(
        val_df,
        x_col='filename',
        y_col='class',
        target_size=(224, 224),
        batch_size=batch_size,
        class_mode='categorical',
        shuffle=False
    )
    
    # 获取类别映射
    class_indices = train_generator.class_indices
    label_to_index = {name: idx for name, idx in class_indices.items()}
    
    return train_generator, val_generator, val_generator, label_to_index

def create_data_generators(images, labels, class_names, test_size=0.2, validation_size=0.2):
    """
    创建训练、验证和测试数据生成器
    """
    # 首先分割训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        images, labels, test_size=test_size, random_state=42, stratify=labels
    )
    
    # 再从训练集中分割出验证集
    X_train, X_val, y_train, y_val = train_test_split(
        X_train, y_train, test_size=validation_size, random_state=42, stratify=y_train
    )
    
    print(f"数据分割完成：")
    print(f"- 训练集: {len(X_train)} 样本")
    print(f"- 验证集: {len(X_val)} 样本")
    print(f"- 测试集: {len(X_test)} 样本")
    
    # 创建标签到索引的映射
    label_to_index = {label: idx for idx, label in enumerate(class_names)}
    
    # 转换标签为one-hot编码
    from tensorflow.keras.utils import to_categorical
    
    y_train_categorical = to_categorical([label_to_index[label] for label in y_train], len(class_names))
    y_val_categorical = to_categorical([label_to_index[label] for label in y_val], len(class_names))
    y_test_categorical = to_categorical([label_to_index[label] for label in y_test], len(class_names))
    
    # 创建数据增强生成器（用于训练集）
    train_datagen = ImageDataGenerator(
        rotation_range=20,      # 随机旋转角度范围
        width_shift_range=0.2,  # 水平平移范围
        height_shift_range=0.2, # 垂直平移范围
        shear_range=0.2,        # 剪切变换范围
        zoom_range=0.2,         # 随机缩放范围
        horizontal_flip=True,   # 水平翻转
        fill_mode='nearest'    # 填充模式
    )
    
    # 验证和测试数据生成器（不进行数据增强）
    val_test_datagen = ImageDataGenerator()
    
    # 创建数据生成器
    batch_size = 32
    
    train_generator = train_datagen.flow(
        X_train, y_train_categorical, 
        batch_size=batch_size,
        shuffle=True
    )
    
    val_generator = val_test_datagen.flow(
        X_val, y_val_categorical,
        batch_size=batch_size,
        shuffle=False
    )
    
    test_generator = val_test_datagen.flow(
        X_test, y_test_categorical,
        batch_size=batch_size,
        shuffle=False
    )
    
    return train_generator, val_generator, test_generator, label_to_index

def visualize_samples(images, labels, class_names, num_samples=5):
    """
    可视化数据集样本
    """
    plt.figure(figsize=(15, 8))
    
    for i, class_name in enumerate(class_names[:min(5, len(class_names))]):
        # 找到该类别的样本索引
        class_indices = np.where(labels == class_name)[0]
        
        if len(class_indices) > 0:
            # 随机选择样本
            sample_indices = np.random.choice(class_indices, min(num_samples, len(class_indices)), replace=False)
            
            for j, idx in enumerate(sample_indices):
                plt_idx = i * num_samples + j + 1
                plt.subplot(len(class_names[:5]), num_samples, plt_idx)
                
                # 显示图像
                plt.imshow(images[idx])
                plt.title(f'{class_name}')
                plt.axis('off')
    
    plt.tight_layout()
    # 保存到docs目录
    os.makedirs('docs', exist_ok=True)
    plt.savefig('docs/dataset_samples.png', dpi=300, bbox_inches='tight')
    plt.show()

def main():
    """主函数"""
    print("=" * 60)
    print("手势识别项目 - 数据预处理模块")
    print("=" * 60)

    # 步骤1：扫描数据目录，创建 DataFrame
    df, class_names = create_dataframe_from_directory("data/leapGestRecog")
    
    if df is None:
        print("数据扫描失败，请检查数据目录")
        return
    
    print(f"\n发现 {len(class_names)} 个手势类别: {class_names}")
    
    # 步骤2：创建数据生成器（从原始位置读取，不移动文件）
    train_gen, val_gen, test_gen, label_map = create_data_generators_from_dataframe(
        df,
        class_names,
        batch_size=32,
        validation_split=0.2
    )
    
    # 保存类别映射到models目录
    import json
    os.makedirs('models', exist_ok=True)
    with open('models/class_mapping.json', 'w') as f:
        json.dump(label_map, f, indent=2)
    
    print("\n" + "=" * 60)
    print("✓ 数据预处理完成！")
    print("=" * 60)
    print("\n生成的文件：")
    print("- models/class_mapping.json: 类别标签映射")
    print("\n优势：")
    print("- 直接从原始位置读取，不移动/复制文件")
    print("- 使用生成器从磁盘读取，无需加载全部数据到内存")
    print("- 支持训练全部 20000 张图像")
    print("\n下一步：构建和训练手势分类模型")
    print("命令: python train_gesture_model.py")

if __name__ == "__main__":
    main()